use rig::agent::{AgentBuilder, AgentBuilderSimple};
use rig::client::completion::CompletionModelHandle;
use rig::providers;
use rig::providers::openrouter::Client;
use rig::{
    agent::PromptRequest,
    client::{completion::CompletionClientDyn, EmbeddingsClient},
    completion::{Chat, Prompt, PromptError},
    embeddings::EmbeddingsBuilder,
    message::{AssistantContent, Message, UserContent},
    vector_store::{
        request::VectorSearchRequestBuilder, VectorSearchRequest, VectorStoreIndex,
        VectorStoreIndexDyn,
    },
    Embed, OneOrMany,
};
use rmcp::{
    model::{ClientCapabilities, ClientInfo, Implementation},
    transport::{ConfigureCommandExt, TokioChildProcess},
    ServiceExt,
};

use crate::util::property;

pub async fn with_mcp_client(
    agent_builder: AgentBuilder<CompletionModelHandle<'_>>,
) -> AgentBuilderSimple<CompletionModelHandle<'_>> {
    let client_info = ClientInfo {
        protocol_version: Default::default(),
        capabilities: ClientCapabilities::default(),
        client_info: Implementation {
            name: "rig-core".to_string(),
            version: "0.13.0".to_string(),
            ..Default::default()
        },
    };
    let transport = TokioChildProcess::new(
        tokio::process::Command::new("/home/x/mcp/java-parser").configure(|cmd| {
            cmd.arg("--mcp");
        }),
    )
    .expect("java-parser create fail");
    let tool_client = client_info
        .serve(transport)
        .await
        .expect("create tool_client fail");
    let tools = tool_client.list_all_tools().await.expect("list tools fail");

    return agent_builder.rmcp_tools(tools, tool_client.peer().to_owned());
}

#[tokio::test]
async fn test_with_mcp_client() {
    // 大模型
    let env: property::Env = property::Env::init();
    let agent_client: Client = providers::openrouter::Client::builder(&env.llm_api_key)
        .base_url(&env.llm_endpoint)
        .build();
    let agent_builder = agent_client.agent(&env.llm_model).temperature(0.5);
    let agent = with_mcp_client(agent_builder).await.build();

    // 提问
    let resp = agent
        .chat(
            "项目路径 /home/x/company/gxdac,有多少接口。如果功能根据查询报错，请完整输出错误"
                .to_string(),
            vec![],
        )
        .await
        .unwrap_or("".to_string());
    println!("resp: {resp}");
}

#[tokio::test]
async fn test_mcp() {
    // 显示mcp的工具集
    let service = ()
        .serve(
            TokioChildProcess::new(
                tokio::process::Command::new("/home/x/mcp/java-parser").configure(|cmd| {
                    cmd.arg("--mcp");
                }),
            )
            .expect("java-parser create fail"),
        )
        .await
        .expect("create service fail");
    let server_info = service.peer_info();
    let tools = service
        .list_tools(Default::default())
        .await
        .expect("list_tools fail")
        .tools;
    println!("Available tools: {tools:#?}");

    let client_info = ClientInfo {
        protocol_version: Default::default(),
        capabilities: ClientCapabilities::default(),
        client_info: Implementation {
            name: "rig-core".to_string(),
            version: "0.13.0".to_string(),
            ..Default::default()
        },
    };
    let transport = TokioChildProcess::new(
        tokio::process::Command::new("/home/x/mcp/java-parser").configure(|cmd| {
            cmd.arg("--mcp");
        }),
    )
    .expect("java-parser create fail");
    let tool_client = client_info
        .serve(transport)
        .await
        .expect("create tool_client fail");

    // 大模型
    let env: property::Env = property::Env::init();
    let agent_client = providers::openrouter::Client::builder(&env.llm_api_key)
        .base_url(&env.llm_endpoint)
        .build();
    let agent = agent_client
        .agent(&env.llm_model)
        .temperature(0.5)
        .preamble("项目下有多少接口？")
        .rmcp_tools(tools, tool_client.peer().to_owned())
        .build();

    // 提问
    let resp = agent
        .chat("项目路径 /home/x/company/gxdac".to_string(), vec![])
        .await
        .unwrap_or("".to_string());
    println!("resp: {resp}");
}
